
import java.io.*;
import java.net.*;
import java.util.*;
import java.nio.*;

class TCPServer extends TCP {
    //Bob RSA key

    private static int e = 3;
    private static int d = 147;
    private static int n = 11 * 23; // p*q

    public static void main(String argv[]) throws Exception
    {
        try {
            ServerSocket serverSocket = new ServerSocket(5555);

            socket = serverSocket.accept();
            in = new DataInputStream(socket.getInputStream());
            out = new DataOutputStream(socket.getOutputStream());

            byte[] encryptionKey = null;
			byte[] integrityKey = null;
			byte[] IV = null;

            while (true) {
				if (socket.isClosed())	
				{
					socket = serverSocket.accept();
					in = new DataInputStream(socket.getInputStream());
					out = new DataOutputStream(socket.getOutputStream());
				}
                //get request from alice
                Message m = ReceiveMessage();

                if (m.type == Message.MessageType.SessionStartRequest) {
                    if (startSession(m)) {
                        System.out.println("Session established successfully");

                        encryptionKey = Utils.HMAC(k, new byte[]{0x01});
                        integrityKey = Utils.HMAC(k, new byte[]{0x02});
                        IV = Arrays.copyOfRange(Utils.HMAC(k, new byte[]{0x03}), 0, 8);
                    } else {
                        sessionId = 0;
						in.close();
						out.close();
						socket.close();
                        System.out.println("Session initiation failed");
                    }
                } else if (m.type == Message.MessageType.SessionResumeRequest) {
                    if (resumeSession(m)) {
                        System.out.println("Session established successfully");
                        encryptionKey = Utils.HMAC(k, new byte[]{0x01});
                        integrityKey = Utils.HMAC(k, new byte[]{0x02});
                        IV = Arrays.copyOfRange(Utils.HMAC(k, new byte[]{0x03}), 0, 8);
                    } else {
						in.close();
						out.close();
						socket.close();
                        System.out.println("Session resume failed");
                    }
                } else if (m.type == Message.MessageType.FileReceive) {
                    if (sessionId != 0) {
                        String fileSendName = Utils.ByteToString(m.data);
                        byte[] fileBytes = Utils.getFileBytes(fileSendName);
                        EncryptClass sendFile = new EncryptClass(cipher.toString(), encryptionKey, integrityKey, IV, fileBytes, true);
						
                        m = new Message(Message.MessageType.FileReceive);
						m.data = sendFile.getCipher();
						
						SendMessage(SimulateAttacker(m));
						
						in.close();
						out.close();
						socket.close();
                    } else {
                        System.out.println("There is no current active session");
                    }
                } else if (m.type == Message.MessageType.FileSend) {
                    if (sessionId != 0) {
                        byte[] cipherBytes = m.data;
                        EncryptClass decryptFile = new EncryptClass(cipher.toString(), encryptionKey, integrityKey, IV, cipherBytes, false);
                        String fileSaveName = Utils.GetUserInput("Enter filename to save as: ");
                        byte[] fileBytes = decryptFile.getText();
                        Utils.writeFile(fileSaveName, fileBytes);
						
						in.close();
						out.close();
						socket.close();
						
                    } else {
                        System.out.println("There is no current active session");
                    }
                } else if (m.type == Message.MessageType.SessionEstablished) {
					//using SessionEstablished as exit message
					System.out.println("Connection with client ended. Server will close now.");
					break;
				}
            }
        } catch (Exception e) {
            System.out.println("Connection with client closed. Server will close now.");
        }
    }

    public static boolean startSession(Message m) throws Exception {
        //values to remember
        byte[] bobR;
        byte[] aliceR;
        byte[] messagesToHash = new byte[56];

        PrintToConsole(false, m);
        aliceR = m.r;
        //save message for hashing later (store in first 8 bytes)
        System.arraycopy(m.ToBytes(), 0, messagesToHash, 0, 8);

        //remember r and sessionId
        bobR = IntToByte(getRandom());
        sessionId = getRandom();
        //send SessionStartResponce to client
        m = new Message(Message.MessageType.SessionStartResponce);
        m.r = bobR;
        m.sessionId = sessionId;
        //set bob's public key (certificate)
        m.e = IntToByte(TCPServer.e);
        m.n = IntToByte(TCPServer.n);
        //promt user to enter ciper selection
		cipher = m.cipher = GetCipherSelection();
        PrintToConsole(true, m);
        //save message for hashing later
        System.arraycopy(m.ToBytes(), 0, messagesToHash, 8, 24);
        SendMessage(SimulateAttacker(m));

        //get PreMasterSecret from client
        m = ReceiveMessage();
        PrintToConsole(false, m);
        //decrypt m.s using Bob's private key
        RsaEncryption rsa = new RsaEncryption(m.s, IntToByte(TCPServer.d), IntToByte(n));
        s = rsa.getCipherArray();
        //now that we have S, aliceR and bobR we can get K = f(S, aliceR, bobR)
        k = GenerateK(s, aliceR, bobR);
        //save message for hashing later
        System.arraycopy(m.ToBytes(), 0, messagesToHash, 32, 24);

        //create keyed hash that we should receive from alice
        byte[] hash = Utils.HMAC(k, Arrays.copyOfRange(messagesToHash, 0, 32));

        //check received hash is correct and return result
        boolean match = Arrays.equals(hash, m.hashOfMsgs);

        if (match) {
            //send KeyedHashOfMsgs to client
            m = new Message(Message.MessageType.KeyedHashOfMsgs);
            //set keyed hash of messages
            m.hashOfMsgs = Utils.HMAC(k, messagesToHash);
            PrintToConsole(true, m);
            SendMessage(SimulateAttacker(m));

            //get Session Established from client
            m = ReceiveMessage();
            PrintToConsole(false, m);
            if (m.type == Message.MessageType.SessionEstablished) {
                return true;
            }
        } else {
            //send error to client
            m = new Message(Message.MessageType.Error);
            m.errorMessage = "Keyed hash of msgs it not correct. Session initiation failed";
            SendMessage(m);
            PrintToConsole(true, m);
        }

        return false;
    }

    public static boolean resumeSession(Message m) throws Exception {
        //values to remember
        byte[] bobR;
        byte[] aliceR;
        byte[] messagesToHash = new byte[44];

        PrintToConsole(false, m);
        if (m.sessionId == sessionId) {
            aliceR = m.r;
            //save message for hashing later (store in first 12 bytes)
            System.arraycopy(m.ToBytes(), 0, messagesToHash, 0, 12);

            //remember r and sessionId
            bobR = IntToByte(getRandom());
            sessionId = getRandom();

            //now that we have S, aliceR and bobR we can get K = f(S, aliceR, bobR)
            k = GenerateK(s, aliceR, bobR);

            //send SessionResumeResponce to client
            m = new Message(Message.MessageType.SessionResumeResponce);
            m.r = bobR;
            m.sessionId = sessionId;
            //promt user to enter ciper selection
           cipher = m.cipher = GetCipherSelection();
            //set keyed hash of messages
            m.hashOfMsgs = Utils.HMAC(k, Arrays.copyOfRange(messagesToHash, 0, 12));
            PrintToConsole(true, m);
            //save message for hashing later
            System.arraycopy(m.ToBytes(), 0, messagesToHash, 12, 32);
            SendMessage(SimulateAttacker(m));

            //get KeyedHashOfMsgs from client
            m = ReceiveMessage();
            PrintToConsole(false, m);
            if (m.type == Message.MessageType.KeyedHashOfMsgs) {
                //check received hash is correct
                //create keyed hash that we should receive from alice
                byte[] hash = Utils.HMAC(k, messagesToHash);
                if (Arrays.equals(hash, m.hashOfMsgs)) {
                    m = new Message(Message.MessageType.SessionEstablished);
                    PrintToConsole(true, m);
                    SendMessage(m);
                    return true;
                }
            }
        } else {
            //send error to client
            m = new Message(Message.MessageType.Error);
            m.errorMessage = "sessionId is not valid";
            SendMessage(m);
            PrintToConsole(true, m);
        }

        return false;

    }

    public static Message.Ciphers GetCipherSelection() throws Exception {
        System.out.println();
        System.out.println("Please make cipher selection:\n");
        System.out.println("1. RC4");
        System.out.println("2. CFB");
        System.out.println("3. CBC");
        System.out.println("4. PCBC");
        String s = Utils.GetUserInput("Enter selection: ");

        if (s.equals("1")) {
            return Message.Ciphers.RC4;
        } else if (s.equals("2")) {
            return Message.Ciphers.CFB;
        } else if (s.equals("3")) {
            return Message.Ciphers.CBC;
        } else {
            return Message.Ciphers.PCBC;
        }
    }
}